
# --*-- coding:utf-8 --*--
"""
    - 能够同时抽取文件中的的文字和图片
    - pdf 和 docx(不支持doc)，可以根据平台进行转换
    - 图片和文本暂时存储到一个临时的文件夹，看师兄如何调用image caption ， 后续应该删除，没有必要存在本地
    - docx支持的图片格式可以通过规则增加
    - 更换了解析pdf的库，但是还没有自己测试速度
"""


from os import path


def filter_document(document):
    """
        todo :add rules
    """
    punctuation = '！“”#¥&：/……～+-\t\n'
    for p in punctuation:
        document = document.replace(p, '')
    return document


def clear_temp(tmp_dir):
    """
         删除临时文件夹使用
    """
    import shutil
    shutil.rmtree(tmp_dir)
    
    # print("Warning rm dir : {} successful!".format(tmp_dir))

def get_image_path(dirs):
    """
        文档里面每张图片的绝对路径以一个列表的方式返回
    """
    import os 
    paths =[]
    for root , _ , files in  os.walk(dirs):
        for file in files:
            paths.append(os.path.join(root,file))
    return paths

def extract_pdf(file, tmp_dir):
    """
        将pdf文件里面的图片保存在临时文件夹之中
    """
    import fitz
    import io
    from PIL import Image
    import os

    if not os.path.exists(tmp_dir):
        os.mkdir(tmp_dir)

    try:
        pdf_file = fitz.open(file)
        for page_index in range(len(pdf_file)):
            # get the page
            page = pdf_file[page_index]
            # image
            for image_index, img in enumerate(page.getImageList(), start=1):
                # get the XREF of the image
                xref = img[0]
                # extract the image bytes
                base_image = pdf_file.extractImage(xref)
                image_bytes = base_image["image"]
                # get the image extension
                image_ext = base_image["ext"]
                image = Image.open(io.BytesIO(image_bytes))
                # save
                save_path = os.path.join(
                    tmp_dir, f"image{page_index+1}_{image_index}.{image_ext}")
                with open(save_path, 'wb') as f:
                    image.save(f)
                f.close()
    except:
        #  如果无法解析pdf则跳过这个文件，后续出现bug再解决
        pass


def extract_doc(docx_path, tmp_dir):
    """
        将docx文件里面的图片保存在临时文件之中
        todo: office和 wps之多格式存在差别
    """
    import docx
    import zipfile
    import os

    if not os.path.exists(tmp_dir):
        os.mkdir(tmp_dir)

    try:
        def extract_images(docx_path, tmp_dir):
            doc = zipfile.ZipFile(docx_path)
            for info in doc.infolist():
                # print(info.filename)/\
                if info.filename.endswith((".png", ".jpeg", ".gif", ".jpg")):
                    doc.extract(info.filename, tmp_dir)
            doc.close()
        extract_images(docx_path, tmp_dir)
    except:
        #  如果无法解析pdf则跳过这个文件，后续出现bug再解决
        pass





# def get_image_path(file_path,dirs='./temp'):
#     """
#         每张图片的绝对路径以一个字符串形式返回，\n 作为分割
#     """
#     import os 
#     paths =[]
#     for root , _ , files in  os.walk(dirs):
#         for file in files:
#             paths.append(os.path.join(root,file))

#     def log():
#         if len(paths)==0:
#             print("{} has no images !".format(file_path))
#         else:
#             print("{} has {} images!".format(file_path , len(paths)))
#     # # log()

#     return '\n'.join(paths)


# method 2
#  python -m fitz extract  -images 'Abst-example/pdf/2021.naacl-main.118.pdf' -output ./test
